import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import os

# 输入模型路径
saved_model_dir = 'd:/Workspace/Python/movie-clip/model/SC'
output_graph = 'd:/Workspace/Python/movie-clip/model/SC_F/frozen_model.pb'

def convert_to_frozen():
    # 加载SavedModel
    model = tf.saved_model.load(saved_model_dir)
    
    # 获取签名或默认计算图
    if not model.signatures:
        print("警告: 模型中没有找到签名，尝试使用默认计算图")
        @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
        def default_func(inputs):
            return model(inputs)
        concrete_func = default_func.get_concrete_function()
    else:
        signature_key = list(model.signatures.keys())[0]
        concrete_func = model.signatures[signature_key]
    
    # 转换为冻结图
    frozen_func = convert_variables_to_constants_v2(concrete_func)
    
    # 确保输出目录存在
    output_dir = os.path.dirname(output_graph)
    if not os.path.exists(output_dir):
        try:
            os.makedirs(output_dir)
            print(f"已创建输出目录: {output_dir}")
        except Exception as e:
            print(f"创建目录失败: {str(e)}")
            # 尝试删除可能存在的文件
            if os.path.exists(output_graph):
                try:
                    os.remove(output_graph)
                    print(f"已删除旧文件: {output_graph}")
                except Exception as e:
                    print(f"删除文件失败: {str(e)}")
            raise
    
    # 保存冻结模型
    try:
        # 使用二进制模式写入
        with open(output_graph, 'wb') as f:
            f.write(frozen_func.graph.as_graph_def().SerializeToString())
        print(f"冻结模型已保存到: {output_graph}")
    except Exception as e:
        print(f"保存模型时出错: {str(e)}")
        raise

if __name__ == '__main__':
    convert_to_frozen()